MODULE fast_cosine_transform

   USE fmm_global_paras
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: cosft2

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE swap(a,b)

      IMPLICIT NONE
      REAL(REALK), INTENT(INOUT) :: a, b
      REAL(REALK) :: tmp
      tmp = a
      a = b
      b = tmp

   END SUBROUTINE

!-------------------------------------------------------------------------------
! Replaces arr(:) by its discrete Fourier transform if flag=1, or replaces
! arr(:) by nn* inverse discrete Fourier transform if flag=-1.
! See Numerical Recipes (edn.2) p.507

   SUBROUTINE four1(nn,flag,arr)
! nn MUST be of form 2^k

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)    :: nn
      INTEGER(INTK), INTENT(IN)    :: flag
      REAL(REALK),   INTENT(INOUT) :: arr(2*nn)

      INTEGER(INTK) :: n, i,j, m, mmax, istep
      REAL(REALK)   :: wtemp, wr,wpr,wpi,wi, theta, tempr,tempi

      n = nn*2
      j = 1
      DO i = 1, n-1, 2
         IF (j > i) THEN
            CALL swap(arr(j),arr(i))
            CALL swap(arr(j+1),arr(i+1))
         END IF
         m = nn
         DO WHILE ( (m >= 2) .AND. (j > m) )
            j = j-m
            m = m/2
         END DO
         j = j+m
      END DO

      mmax = 2
      DO WHILE (n > mmax)
         istep = mmax *2
         theta = flag * (6.28318530717959d0 / mmax)
         wtemp = SIN(half*theta)
         wpr = -two*wtemp*wtemp
         wpi = SIN(theta)
         wr = one
         wi = zero

         DO m = 1, mmax-1, 2
            DO i = m, n, istep
               j = i+mmax
               tempr = wr*arr(j)   -wi*arr(j+1)
               tempi = wr*arr(j+1) +wi*arr(j)
               arr(j) = arr(i) - tempr
               arr(j+1) = arr(i+1) - tempi
               arr(i) = arr(i) + tempr
               arr(i+1) = arr(i+1) + tempi
            END DO
            wtemp = wr
            wr = wr*wpr -wi*wpi +wr
            wi = wi*wpr +wtemp*wpi +wi
         END DO
         mmax = istep

      END DO

   END SUBROUTINE four1

!-------------------------------------------------------------------------------
! Calculates the Fourier transform of a set of n real-valued data points.
! See Numerical Recipes (edn.2) p.513

   SUBROUTINE realft(n,flag,arr)
! n MUST be of form 2^k

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)    :: n
      INTEGER(INTK), INTENT(IN)    :: flag
      REAL(REALK),   INTENT(INOUT) :: arr(n)

      INTEGER(INTK) :: i, i1,i2,i3,i4, np3
      REAL(REALK)   :: wtemp, wr,wpr,wpi,wi, theta
      REAL(REALK)   :: c1,c2, h1r,h2r, h1i,h2i

      theta = PI / (n/2)

      c1 = half
      IF (flag == 1) THEN
         c2 = -half
         CALL four1(n/2,1,arr)
      ELSE
         c2 = half
         theta = -theta
      END IF

      wtemp = SIN(half*theta)
      wpr = -two*wtemp*wtemp
      wpi = SIN(theta)
      wr = one + wpr
      wi = wpi
      np3 = n + 3

      DO i = 2, n/4
         i1 = i +i -1
         i2 = 1 +i1
         i3 = np3 -i2
         i4 = 1 +i3
         h1r = c1*(arr(i1)+arr(i3))
         h1i = c1*(arr(i2)-arr(i4))
         h2r = -c2*(arr(i2)+arr(i4))
         h2i = c2*(arr(i1)-arr(i3))
         arr(i1) = h1r +wr*h2r -wi*h2i
         arr(i2) = h1i +wr*h2i +wi*h2r
         arr(i3) = h1r -wr*h2r +wi*h2i
         arr(i4) = -h1i +wr*h2i +wi*h2r
         wtemp = wr
         wr = wr*wpr -wi*wpi +wr
         wi = wi*wpr +wtemp*wpi +wi
      END DO

      IF (flag == 1) THEN
         h1r = arr(1)
         arr(1) = h1r + arr(2)
         arr(2) = h1r - arr(2)
      ELSE
         h1r = arr(1)
         arr(1) = c1*(h1r + arr(2))
         arr(2) = c1*(h1r - arr(2))
         CALL four1(n/2,-1,arr)
      END IF

   END SUBROUTINE realft

!-------------------------------------------------------------------------------
! Calculates the "staggered" cosine transform of a set of y(1:n)
! real-valued data points.
! See Numerical Recipes (edn.2) p.520

   SUBROUTINE cosft2(n,flag,y)
! n MUST be of form 2^k

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)    :: n
      INTEGER(INTK), INTENT(IN)    :: flag
      REAL(REALK),   INTENT(INOUT) :: y(n)

      INTEGER(INTK) :: i
      REAL(REALK)   :: sm,sm1, y1,y2, ytemp
      REAL(REALK)   :: wtemp, wr,wpr,wpi,wi,wi1,wr1, theta

      wi = zero
      wr = one
      theta = half * PI / n
      wr1 = COS(theta)
      wi1 = SIN(theta)
      wpr = -two*wi1*wi1
      wpi = SIN(two*theta)

      IF ( flag == 1 ) THEN  ! Forward transform

         DO i = 1, n/2
            y1 = half*(y(i)+y(n-i+1))
            y2 = wi1 *(y(i)-y(n-i+1))
            y(i) = y1 + y2
            y(n-i+1) = y1 - y2
            wtemp = wr1
            wr1 = wr1*wpr -wi1*wpi +wr1
            wi1 = wi1*wpr +wtemp*wpi +wi1
         END DO

         CALL realft(n,1,y)

         DO i = 3, n, 2
            wtemp = wr
            wr = wr*wpr -wi*wpi +wr
            wi = wi*wpr +wtemp*wpi +wi
            y1 = y(i)*wr -y(i+1)*wi
            y2 = y(i+1)*wr +y(i)*wi
            y(i) = y1
            y(i+1) = y2
         END DO

         sm = half*y(2)
         DO i = n, 2, -2
            sm1 = sm
            sm = sm + y(i)
            y(i) = sm1
         END DO

      ELSE
         STOP 'inverse cosine transform NYI'
      END IF

   END SUBROUTINE cosft2

!-------------------------------------------------------------------------------

END MODULE fast_cosine_transform

